package com.doubleview.sandbox;

import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import groovy.lang.GroovyClassLoader;
import groovy.lang.Script;
import groovy.transform.ConditionalInterrupt;
import groovy.transform.ThreadInterrupt;
import org.codehaus.groovy.ast.ClassHelper;
import org.codehaus.groovy.ast.Parameter;
import org.codehaus.groovy.ast.expr.ClassExpression;
import org.codehaus.groovy.ast.expr.ClosureExpression;
import org.codehaus.groovy.ast.expr.MethodCallExpression;
import org.codehaus.groovy.ast.stmt.ExpressionStatement;
import org.codehaus.groovy.control.CompilerConfiguration;
import org.codehaus.groovy.control.customizers.ASTTransformationCustomizer;
import org.junit.Test;
import org.kohsuke.groovy.sandbox.SandboxTransformer;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.*;
import java.util.function.Supplier;

/**
 * @author huchengchao <huchengchao@kuaishou.com>
 * Created on 2022-03-14
 */
public class GroovyShellTest {


    ListeningExecutorService executor = MoreExecutors.listeningDecorator(
            Executors.newScheduledThreadPool(1));

    private String SCRIPT_STR = "System.exit(0)";

    @Test
    public void scriptTest() throws Exception {
        // 脚本名称,key
        String scriptId = "test";
        // 缓存脚本，提高执行性能

        ListenableFuture future = executor.submit(new Callable<Object>() {
            @Override
            public Object call() throws Exception {
                GroovyClassLoader groovyClassLoader = createClassLoaderSync();
                Script script = (Script) groovyClassLoader.parseClass(SCRIPT_STR).getDeclaredConstructor().newInstance();
                try {
                    return script.run();
                } catch (Exception e) {
                    e.printStackTrace();
                    script.run();
                }
                return null;
            }
        });


        try {
            future.get(5000, TimeUnit.MILLISECONDS);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }


    private GroovyClassLoader createClassLoaderSync() {
        // 自定义配置
        CompilerConfiguration config = new CompilerConfiguration();
        // 添加线程中断拦截器，可拦截循环体（for,while）、方法和闭包的首指令
        config.addCompilationCustomizers(new ASTTransformationCustomizer(ThreadInterrupt.class));
        Map<String, Object> timeoutArgs = new HashMap<>();
        timeoutArgs.put("value", new ClosureExpression(
                Parameter.EMPTY_ARRAY,
                new ExpressionStatement(
                        new MethodCallExpression(
                                new ClassExpression(ClassHelper.make(ConditionInterceptor.class)),
                                "checkTimeout",
                                MethodCallExpression.NO_ARGUMENTS)
                )
        ));
        config.addCompilationCustomizers(new ASTTransformationCustomizer(timeoutArgs, ConditionalInterrupt.class));

        // 沙盒环境
        config.addCompilationCustomizers(new SandboxTransformer());
        // 注册方法拦截
        new GroovyNotSupportInterceptor().register();
        return new GroovyClassLoader(Thread.currentThread().getContextClassLoader(), config, true);
    }


    private GroovyClassLoader createClassLoaderAsync() {
        try {
            return CompletableFuture.supplyAsync(new Supplier<GroovyClassLoader>() {
                @Override
                public GroovyClassLoader get() {
                    // 自定义配置
                    CompilerConfiguration config = new CompilerConfiguration();
                    // 添加线程中断拦截器，可拦截循环体（for,while）、方法和闭包的首指令
                    config.addCompilationCustomizers(new ASTTransformationCustomizer(ThreadInterrupt.class));
                    Map<String, Object> timeoutArgs = new HashMap<>();
                    timeoutArgs.put("value", new ClosureExpression(
                            Parameter.EMPTY_ARRAY,
                            new ExpressionStatement(
                                    new MethodCallExpression(
                                            new ClassExpression(ClassHelper.make(ConditionInterceptor.class)),
                                            "checkTimeout",
                                            MethodCallExpression.NO_ARGUMENTS)
                            )
                    ));
                    config.addCompilationCustomizers(new ASTTransformationCustomizer(timeoutArgs, ConditionalInterrupt.class));

                    // 沙盒环境
                    config.addCompilationCustomizers(new SandboxTransformer());
                    // 注册方法拦截
                    new GroovyNotSupportInterceptor().register();
                    return new GroovyClassLoader(Thread.currentThread().getContextClassLoader(), config, true);
                }
            }).get();
        } catch (InterruptedException e) {
            e.printStackTrace();
        } catch (ExecutionException e) {
            e.printStackTrace();
        }
        return null;
    }
}
